# -*- coding: utf-8 -*-
'''
Created on 26.07.2019

@author: yu03
'''

import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_boxcar, FFT_interpolation_2, FFT_interpolation_compare, FFT_cal
import numpy as np
import os
from scipy.optimize import curve_fit
from scipy import signal


def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return np.concatenate(( np.zeros(int((N)/2)), (cumsum[N:]-cumsum[:-N])/float(N) , np.zeros(int((N-1)/2)) ))
def fit_func(x, a, b, c):
    return a*(x-b)**2 + c
# p= -0.0213125
p= -0.02135
def fit_func_phase(x, q):
    return p*x+q

def fit_func_line(x, p, q):
    return p*x+q



#############
# fs_cam = 575.8  ###USB2.0
# fs_cam = 1370.72 ###USB3.0
fs_cam = 1000 ### Hardware Trigger

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_phase_centers = []
hor_freq_set = []
hor_f_fit_set = []
hor_phi_fit_set = []

file_path = r'C:\Users\yu03\Desktop\Experiment Record\!!! Isolation Room Test\Upstair well alighned\20190808 nonliearity test\Black Head\Noise_PSD\New Isolation'
doc_name = '50Hz_closed_mode'
file_type = '.avi'
file_name = file_path + '\\' +  doc_name + file_type

cap = cv2.VideoCapture(file_name)
frame_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
print(frame_num, "Frames")


for k in range(frame_num):
    ret, frame = cap.read() 
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    img_set.append(img)
print('Video reading ok')

num = 0
for img in img_set[:]:
    num += 1
    hor_center = img[0]
    DC_num = 1000
    hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
    hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][0]+DC_num
    hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][-1]+DC_num
    hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
    hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]
    hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
    hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
    [hor_a, hor_b, hor_c] = hor_params[0]
    hor_f_fit = hor_b
    hor_f_fit_set.append(hor_f_fit)
    hor_params = curve_fit(fit_func_phase, hor_fit_x, np.unwrap(hor_fit_phase))
    [hor_q] = hor_params[0]
    hor_phi_fit = p*hor_freqline[hor_m_k_num]+hor_q
    hor_phi_fit_set.append(hor_phi_fit+2*np.pi-0.065)
    print(num, hor_m_k_num)
    
    hor_freq_set.append(hor_freq_estim)    
    hor_phase_centers.append(hor_phase_estim)


############### linearity range
start_cutting, stop_cutting = 0, len(hor_phase_centers)
# start_cutting, stop_cutting = 0, 4500
selected_range = np.linspace(start_cutting, stop_cutting, stop_cutting-start_cutting)

plt.figure(1)
plt.subplot(2,1,1)
plt.plot(hor_center)
plt.grid(which='major', axis='both')
plt.subplot(2,1,2)
#     plt.stem(hor_freqline, hor_sig_magnitude, use_line_collection=True)
plt.plot(hor_freqline, hor_sig_magnitude)
plt.grid(which='major', axis='both')
plt.show()


fig = plt.figure(2)
plt.gcf().set_size_inches(18,9)
ax1 = plt.subplot(4,2,1)
plt.plot(hor_f_fit_set, color='blue', marker=' ')
measured_freq = hor_f_fit_set[start_cutting:stop_cutting]
fit_params = curve_fit(fit_func_line, np.linspace(0, len(measured_freq)-1, len(measured_freq)), measured_freq)
# [p, q] = fit_params[0]
[p, q] = np.polyfit(np.linspace(0, len(measured_freq)-1, len(measured_freq)), measured_freq, 1)
linear_fit = np.linspace(0, len(measured_freq)-1, len(measured_freq))*p + q
hor_freq_nonlinearity = (measured_freq - linear_fit)
plt.plot(selected_range, linear_fit, color='red')
plt.title("Horizontal FFT")
plt.ylabel("Horizontal Freq Estimation (/m)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
ax1_angle = ax1.twinx()
plt.ylim(angle_range)
plt.ylabel('Horizontal Tilting (urad)')
  
ax2 = plt.subplot(4,2,2)
hor_phase_unwrap = np.unwrap(hor_phase_centers)
plt.plot(hor_phase_unwrap, color='blue', marker=' ')
# plt.plot(hor_phi_fit_set, color='red', marker=' ')
measured_phase = hor_phase_unwrap[start_cutting:stop_cutting]
fit_params = curve_fit(fit_func_line, np.linspace(0, len(measured_phase)-1, len(measured_phase)), np.unwrap(measured_phase))
# [p, q] = fit_params[0]
[p, q] = np.polyfit(np.linspace(0, len(measured_phase)-1, len(measured_phase)), np.unwrap(measured_phase), 1)
linear_fit = np.linspace(0, len(measured_phase)-1, len(measured_phase))*p + q
hor_phase_nonlinearity = (measured_phase - linear_fit)
plt.plot(selected_range, linear_fit, color='red')
plt.title("Horizontal Phase")
plt.ylabel("Horizontal Phase (rad)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax2_length = ax2.twinx()
plt.ylim(length_range)
plt.ylabel('Horizontal Length (nm)')

ax3 = plt.subplot(4,2,3)
plt.plot(selected_range, hor_freq_nonlinearity, color='blue')
plt.title("Horizontal Tilting Noise_PSD")
plt.ylabel("Horizontal Freq Estimation (/m)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
ax3_angle = ax3.twinx()
plt.ylim(angle_range)
plt.ylabel('Horizontal Tilting (urad)')
 
ax4 = plt.subplot(4,2,4)
plt.plot(selected_range, hor_phase_nonlinearity, color='blue')
plt.title("Horizontal Phase Noise_PSD")
plt.ylabel("Horizontal Phase (rad)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax4_length = ax4.twinx()
plt.ylim(length_range)
plt.ylabel('Horizontal Length (nm)')

ax5 = plt.subplot(4,2,5)
plt.plot(selected_range, hor_freq_nonlinearity, color='blue')
plt.title("Horizontal Tilting Noise_PSD")
plt.ylabel("Horizontal Freq Estimation (/m)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
ax5_angle = ax5.twinx()
plt.ylim(angle_range)
plt.ylabel('Horizontal Tilting (urad)')

ax6 = plt.subplot(4,2,6)
plt.plot(selected_range, hor_phase_nonlinearity, color='blue')
plt.title("Horizontal Phase Noise_PSD")
plt.ylabel("Horizontal Phase (rad)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax6_length = ax6.twinx()
plt.ylim(length_range)
plt.ylabel('Horizontal Length (nm)')

# plt.tight_layout(pad=0.4, w_pad=0.5, h_pad=1.0)

# freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(hor_f_fit_set, 1/fs_cam)
# plt.figure('Noise_PSD Spectrum')
ax1 = plt.subplot(4,2,7)
# freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(hor_f_fit_set, 1/fs_cam)
f_line, PS = signal.welch(hor_f_fit_set, fs_cam, nperseg=len(hor_f_fit_set))
PSD = np.sqrt(PS)
plt.semilogy(f_line, PSD, 'b')
plt.title("Freq. Noise_PSD FFT")
plt.ylabel("Freq. Noise_PSD (/m/sqrt(Hz)")
plt.xlabel("Freq. (Hz)")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
ax1_angle = ax1.twinx()
plt.ylim(angle_range)
plt.ylabel('Angle Noise_PSD (urad/sqrt(Hz)')
ax2 = plt.subplot(4,2,8)
# freqline, sig_FFT, sig_magnitude, sig_phase = FFT_cal(hor_phase_unwrap, 1/fs_cam)
f_line, PS = signal.welch(hor_phase_unwrap, fs_cam, nperseg=len(hor_phase_unwrap))
PSD = np.sqrt(PS)
plt.semilogy(f_line, PS, 'b')
plt.title("Phase Noise_PSD FFT")
plt.ylabel("Phase Noise_PSD(rad/sqrt(Hz))")
plt.xlabel("Freq. (Hz)")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax2_length = ax2.twinx()
plt.ylim(length_range)
plt.ylabel('Length Noise_PSD (nm/sqrt(Hz))')
fig.tight_layout()
plt.show()